'''
Created on Nov 20, 2013

@author: maxp
'''

import sys
import argparse
from Bio import SeqIO as bseq
from numpy import arange, zeros
import align
from progressbar import ETA, ProgressBar, SimpleProgress, Percentage, Bar, FileTransferSpeed



def main():
    parser = argparse.ArgumentParser(description="For each read in 'reference' find all reads in 'query' within an edit distance")
    parser.add_argument('ref_file', nargs='?', default="reference.fna")
    parser.add_argument('query_file', nargs='?', default="query.fna")
    parser.add_argument('--out_file', nargs='?', const="ref_matches.txt", type=argparse.FileType('w'),
                        default=sys.stdout)
    parser.add_argument('--maxedit', default=6, type=int)
    
    args = parser.parse_args()
        
    #longest query and longest reference both 634 characters
    T = zeros(shape=(635,635),dtype=int)
    T[0] = arange(635)
    T[:,0] = arange(635)
    qs_and_ids = sorted((str(rec.seq),rec.id) for rec in bseq.parse(args.query_file,"fasta"))
        
    widgets = ['Progress: ', SimpleProgress(), '-', Percentage(), ' ',
               Bar(marker='0',left='[',right=']'), ETA(), FileTransferSpeed()]
                                                                      
        
    for ref_rec in bseq.parse(args.ref_file,"fasta"):
        ref_seq = str(ref_rec.seq)
        result = [ref_rec.id]
        qnum = 0
        pbar = ProgressBar(widgets=widgets, maxval=len(qs_and_ids)).start()

        while qnum < len(qs_and_ids):
            qseq, qid = qs_and_ids[qnum]
            
            if len(qseq) > len(ref_seq) + args.maxedit:
                aligned  = align.align(qseq[:len(ref_seq)+ args.maxedit],
                                                ref_seq, 
                                                T, 1, args.maxedit)
            elif len(ref_seq) > len(qseq) + args.maxedit:
                aligned  = align.align(qseq,
                                                ref_seq[:len(qseq)+args.maxedit],
                                                T, 1, args.maxedit)
            else:
                aligned = align.align(qseq,
                                                ref_seq, 
                                                T, 1, args.maxedit)
            if aligned[0]: 
                result.append(qid)
            
            qnum += 1
            if qnum % 5000 == 0: pbar.update(qnum)
            
            
        pbar.finish()        
        args.out_file.write('\t'.join(result))
        args.out_file.write('\n')
        del result
                  
    args.out_file.close()        


if __name__ == '__main__':
    main()